import matplotlib.pyplot as plt
import numpy as np
import argparse


def plot_data(log_file, output_image, xpu_type='GPU'):
    timestamps = []
    xpu_utils = []
    mem_used = []
    mem_total = []

    with open(log_file, 'r') as f:
        for line in f:
            parts = line.strip().split(',')
            if len(parts) == 4:
                timestamps.append(parts[0])
                xpu_utils.append(float(parts[1].strip()))
                mem_used.append(float(parts[2].strip()))
                mem_total.append(float(parts[3].strip()))

    x = np.arange(len(timestamps))

    fig, ax1 = plt.subplots(figsize=(14, 10))

    color = 'tab:blue'
    ax1.set_xlabel('Time')
    ax1.set_ylabel(f'{xpu_type} Utilization (%)', color=color)
    ax1.plot(x, xpu_utils, color=color, label=f'{xpu_type} Utilization')
    ax1.tick_params(axis='y', labelcolor=color)
    ax1.set_xticks(x[::max(1, len(x)//10)])
    ax1.set_xticklabels(timestamps[::max(1, len(timestamps)//10)], rotation=45, ha='right')

    ax2 = ax1.twinx()
    color = 'tab:green'
    ax2.set_ylabel('Memory Used (MiB)', color=color)
    ax2.plot(x, mem_used, color=color, label='Memory Used')
    ax2.tick_params(axis='y', labelcolor=color)

    fig.tight_layout()
    plt.title(f'{xpu_type} Memory and Utilization Over Time')
    plt.savefig(output_image)
    plt.close()

    # summary result 
    print(f"************************** log_file : {log_file} **************************")
    print(f"[INFO] timestamp num : {len(timestamps)}")
    print(f"[INFO] max {xpu_type} utils : {max(xpu_utils)}")
    print(f"[INFO] max mem_used : {max(mem_used)}")
    print(f"[INFO] max mem_total : {max(mem_total)}")
    print(f"***************************************************************************")

def parse_args():
    """Parse command-line arguments."""
    parser = argparse.ArgumentParser(description='Plot utilization and memory from log files.')
    parser.add_argument('log_files', nargs='+', help='List of log file names to process')
    parser.add_argument('--output', default='output.png', help='Output image file name')
    return parser.parse_args()



